{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9bfe1950",
   "metadata": {},
   "source": [
    "# Text2SQL"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fef6a9ba",
   "metadata": {},
   "source": [
    "Text2SQL的流程：\n",
    "1. 通过LangChain接入AI大模型，构建AI问答服务\n",
    "2. 读取数据库的schema，并构建为Prompt\n",
    "3. 将构建好的包含schema的prompt传给llm获取生成的sql\n",
    "4. 执行SQL来获取结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "179bda28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pymysql\n",
    "from urllib.parse import urlparse\n",
    "from langchain_community.chat_models.tongyi import ChatTongyi\n",
    "from langchain_core.prompts import ChatPromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "779ba441",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatTongyi(\n",
    "    model='qwen-plus',\n",
    "    api_key=\"sk-da90821cf9174fbeb854011015c67aad\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "695cf85a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_sql_query(query, db_path=\"root:123456@localhost:3306/sim_war_game\"):\n",
    "    \"\"\"\n",
    "    执行传入的SQL查询（适用于 MySQL）\n",
    "    args:\n",
    "    - query: 要执行的SQL\n",
    "    - db_path: MySQL连接路径\n",
    "    return：查询的结果（列表 of 字典）\n",
    "    \"\"\"\n",
    "    # 解析 db_path，假设格式为 user:password@host:port/database\n",
    "    try:\n",
    "        # 支持格式：user:pass@host:port/dbname\n",
    "        parsed = urlparse(f'mysql://{db_path}')\n",
    "        netloc = parsed.netloc\n",
    "        path = parsed.path.lstrip('/')\n",
    "\n",
    "        # 解析用户名密码\n",
    "        if '@' in netloc:\n",
    "            auth, host_port = netloc.split('@')\n",
    "            if ':' in auth:\n",
    "                user, password = auth.split(':')\n",
    "            else:\n",
    "                user, password = auth, ''\n",
    "        else:\n",
    "            host_port = netloc\n",
    "            user, password = 'root', ''  # 默认用户\n",
    "\n",
    "        # 解析 host 和 port\n",
    "        if ':' in host_port:\n",
    "            host, port = host_port.split(':')\n",
    "            port = int(port)\n",
    "        else:\n",
    "            host, port = host_port, 3306\n",
    "\n",
    "        database = path or 'mysql'\n",
    "\n",
    "    except Exception as e:\n",
    "        raise ValueError(f\"无法解析数据库路径: {db_path}\") from e\n",
    "\n",
    "    # 连接数据库并执行查询\n",
    "    connection = None\n",
    "    try:\n",
    "        connection = pymysql.connect(\n",
    "            host=host,\n",
    "            port=port,\n",
    "            user=user,\n",
    "            password=password,\n",
    "            database=database,\n",
    "            charset='utf8mb4',\n",
    "            cursorclass=pymysql.cursors.DictCursor  # 返回字典格式\n",
    "        )\n",
    "        with connection.cursor() as cursor:\n",
    "            cursor.execute(query)\n",
    "            result = cursor.fetchall()\n",
    "        connection.commit()\n",
    "        return result\n",
    "\n",
    "    except Exception as e:\n",
    "        raise RuntimeError(f\"SQL执行失败: {e}\") from e\n",
    "\n",
    "    finally:\n",
    "        if connection:\n",
    "            connection.close()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0183cdd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'scene_id': 'a1b2c3d4-e5f6-7890-g1h2-i3j4k5l6m7n8', 'scene_name': '边境区域常态化巡逻', 'scene_type': 'area_patrol', 'red_config': '{\"units\": [\"uav\", \"patrol_vehicle\"], \"duration\": 120, \"patrol_route\": \"route_a\"}', 'blue_config': '{\"units\": [\"observation_post\", \"rapid_response_team\"], \"alert_level\": 2}', 'create_time': datetime.datetime(2025, 9, 11, 7, 30)}]\n"
     ]
    }
   ],
   "source": [
    "query = \"SELECT * FROM t_scene LIMIT 1;\"\n",
    "result = run_sql_query(query, \"root:123456@localhost:3306/sim_war_game\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e8e6d280",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_schema_prompt(db_path=\"root:123456@localhost:3306/sim_war_game\"):\n",
    "    \"\"\"\n",
    "    从 MySQL 数据库路径中提取数据库模式信息。\n",
    "\n",
    "    args:\n",
    "    - db_path: MySQL 连接路径\n",
    "\n",
    "    return:\n",
    "    - 包含所有表、字段、约束及注释的格式化字符串\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 构造伪URL以便使用urlparse\n",
    "        if not db_path.startswith(\"mysql://\"):\n",
    "            db_path = f\"mysql://{db_path}\"\n",
    "        \n",
    "        parsed = urlparse(db_path)\n",
    "        \n",
    "        username = parsed.username\n",
    "        password = parsed.password\n",
    "        host = parsed.hostname\n",
    "        port = parsed.port or 3306\n",
    "        database = parsed.path.lstrip('/')\n",
    "        if not database:\n",
    "            raise ValueError(\"缺少数据库名\")\n",
    "\n",
    "        # 建立数据库连接\n",
    "        connection = pymysql.connect(\n",
    "            host=host,\n",
    "            port=port,\n",
    "            user=username,\n",
    "            password=password,\n",
    "            database=database,\n",
    "            charset='utf8mb4'\n",
    "        )\n",
    "        cursor = connection.cursor()\n",
    "\n",
    "        prompt_lines = []\n",
    "\n",
    "        # 获取所有表名及其注释\n",
    "        cursor.execute(\"\"\"\n",
    "            SELECT TABLE_NAME, TABLE_COMMENT \n",
    "            FROM INFORMATION_SCHEMA.TABLES \n",
    "            WHERE TABLE_SCHEMA = %s AND TABLE_TYPE = 'BASE TABLE';\n",
    "        \"\"\", (database,))\n",
    "        tables = cursor.fetchall()\n",
    "\n",
    "        for table_name, table_comment in tables:\n",
    "            # 添加表名和表注释\n",
    "            prompt_lines.append(f\"Table: {table_name}\")\n",
    "            if table_comment and table_comment.strip():\n",
    "                prompt_lines.append(f\"  Comment: {table_comment.strip()}\")\n",
    "\n",
    "            # 获取列信息（含注释）\n",
    "            cursor.execute(\"\"\"\n",
    "                SELECT \n",
    "                    COLUMN_NAME, \n",
    "                    DATA_TYPE, \n",
    "                    IS_NULLABLE, \n",
    "                    COLUMN_DEFAULT, \n",
    "                    COLUMN_KEY,\n",
    "                    COLUMN_COMMENT\n",
    "                FROM INFORMATION_SCHEMA.COLUMNS \n",
    "                WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s\n",
    "                ORDER BY ORDINAL_POSITION;\n",
    "            \"\"\", (database, table_name))\n",
    "            \n",
    "            columns = cursor.fetchall()\n",
    "            for col in columns:\n",
    "                name, col_type, is_nullable, dflt_value, column_key, col_comment = col\n",
    "                desc = f\" - {name}: {col_type}\"\n",
    "                if is_nullable == 'NO':\n",
    "                    desc += \" NOT NULL\"\n",
    "                if dflt_value is not None:\n",
    "                    # 转义单引号\n",
    "                    if isinstance(dflt_value, str):\n",
    "                        dflt_value = dflt_value.replace(\"'\", \"''\")\n",
    "                    desc += f\" DEFAULT '{dflt_value}'\"\n",
    "                if column_key == 'PRI':\n",
    "                    desc += \" PRIMARY KEY\"\n",
    "                if column_key == 'MUL':\n",
    "                    desc += \" INDEX\"  # 简单标记外键相关索引\n",
    "                prompt_lines.append(desc)\n",
    "\n",
    "                # 添加字段注释（如果存在）\n",
    "                if col_comment and col_comment.strip():\n",
    "                    prompt_lines.append(f\"   > {col_comment.strip()}\")\n",
    "\n",
    "            prompt_lines.append(\"\")  # 表之间空行分隔\n",
    "\n",
    "        cursor.close()\n",
    "        connection.close()\n",
    "\n",
    "        return \"\\n\".join(prompt_lines)\n",
    "\n",
    "    except Exception as e:\n",
    "        raise RuntimeError(f\"提取数据库模式失败: {e}\") from e\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1f3f07cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Table: simulation_steps\n",
      "  Comment: 存储每次模拟任务的各个步骤详情，按步保存状态、动作和事件\n",
      " - id: int NOT NULL PRIMARY KEY\n",
      "   > 主键，自增ID\n",
      " - simulation_id: int NOT NULL INDEX\n",
      "   > 关联的模拟任务ID，外键指向 simulations 表的 id\n",
      " - step_number: int NOT NULL\n",
      "   > 当前步骤的序号，从1开始递增\n",
      " - step_data: json\n",
      "   > 该步骤的状态数据，JSON格式，存储模拟状态快照\n",
      " - action_log: json\n",
      "   > 该步骤执行的动作日志，JSON格式，记录用户或系统操作\n",
      " - events: json\n",
      "   > 该步骤触发的事件列表，JSON格式，如告警、通知等\n",
      " - step_time: datetime NOT NULL\n",
      "   > 该步骤对应的时间点（模拟时间）\n",
      " - created_at: datetime DEFAULT 'CURRENT_TIMESTAMP'\n",
      "   > 记录创建时间，自动填充为当前时间\n",
      "\n",
      "Table: simulations\n",
      "  Comment: 模拟任务主表，记录每次模拟的元信息、状态和进度\n",
      " - id: int NOT NULL PRIMARY KEY\n",
      "   > 主键，自增ID，唯一标识一次模拟任务\n",
      " - name: varchar NOT NULL\n",
      "   > 模拟任务的名称，便于用户识别\n",
      " - scene_id: varchar NOT NULL INDEX\n",
      "   > 关联的场景ID（UUID格式），用于标识所属的模拟场景\n",
      " - status: varchar\n",
      "   > 当前模拟状态，如 pending, running, paused, completed, failed 等\n",
      " - current_step: int\n",
      "   > 当前已执行到的步骤编号\n",
      " - total_steps: int\n",
      "   > 预计或最终的总步骤数\n",
      " - speed_multiplier: int\n",
      "   > 模拟速度倍率，如 1, 2, 5 表示1倍、2倍、5倍速\n",
      " - auto_mode: tinyint\n",
      "   > 是否为自动模式（0: 手动, 1: 自动），控制模拟是否自动推进\n",
      " - initial_state: json\n",
      "   > 模拟初始状态快照，JSON格式，用于重置或回放\n",
      " - current_state: json\n",
      "   > 当前状态快照，JSON格式，用于恢复断点\n",
      " - result_data: json\n",
      "   > 模拟完成后的结果数据，如统计、指标、报告等，JSON格式\n",
      " - start_time: datetime\n",
      "   > 模拟实际开始时间\n",
      " - end_time: datetime\n",
      "   > 模拟结束时间（成功或失败）\n",
      " - created_at: datetime DEFAULT 'CURRENT_TIMESTAMP'\n",
      "   > 记录创建时间，自动填充\n",
      " - updated_at: datetime DEFAULT 'CURRENT_TIMESTAMP'\n",
      "   > 最后更新时间，通常由触发器或应用逻辑更新\n",
      "\n",
      "Table: t_scene\n",
      "  Comment: 模拟场景定义表，存储不同推演场景的配置与类型\n",
      " - scene_id: varchar NOT NULL PRIMARY KEY\n",
      "   > 场景唯一标识符，通常使用UUID\n",
      " - scene_name: varchar NOT NULL\n",
      "   > 场景名称，用于展示和识别\n",
      " - scene_type: enum NOT NULL\n",
      "   > 场景类型：目标打击、区域巡逻、区域控制\n",
      " - red_config: json NOT NULL\n",
      "   > 红方（攻击方）配置信息，JSON格式，包含兵力、策略等\n",
      " - blue_config: json NOT NULL\n",
      "   > 蓝方（防御方）配置信息，JSON格式，包含部署、响应逻辑等\n",
      " - create_time: datetime DEFAULT 'CURRENT_TIMESTAMP'\n",
      "   > 场景创建时间，自动填充为当前时间戳\n",
      "\n"
     ]
    }
   ],
   "source": [
    "db_path = \"root:123456@localhost:3306/sim_war_game\"\n",
    "schema_prompt = extract_schema_prompt(db_path)\n",
    "\n",
    "print(schema_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fa84e99",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_system_prompt(prompt, llm):\n",
    "    \"\"\"\n",
    "    生成系统提示并调用语言模型生成SQL查询。\n",
    "\n",
    "    args:\n",
    "    - prompt: 用户输入的查询需求。\n",
    "    - llm: 用于生成SQL查询的语言模型实例。\n",
    "\n",
    "    return:\n",
    "    - 生成的SQL查询语句。\n",
    "    \"\"\"\n",
    "\n",
    "    # 定义提示模板\n",
    "    prompt_template = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", \"\"\"\n",
    "你是一个SQL专家，请根据提供的表结构信息和查询需求，来生成可执行的mysql支持的sql查询语句，你只能回复可执行的SQL内容，\n",
    "生成的sql放在<SQL>标签内部\n",
    "\n",
    "A：查询scene_id='a1b2c3d4-e5f6-7890-g1h2-i3j4k5l6m7n8'红方配置信息\n",
    "Q：<SQL>select red_config from t_scene where scene_id=\"a1b2c3d4-e5f6-7890-g1h2-i3j4k5l6m7n8\"</SQL>\n",
    "\"\"\"),\n",
    "            (\"human\", \"\"\"\n",
    "【数据库的结构】\n",
    "{schema}\n",
    "\n",
    "【查询需求】\n",
    "{input}\n",
    "\"\"\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # 用于获取数据库结构信息\n",
    "    schema = extract_schema_prompt()\n",
    "\n",
    "    # 使用提示模板生成提示词\n",
    "    prompt_value = prompt_template.invoke({\"input\": prompt, \"schema\": schema})\n",
    "\n",
    "    # 调用语言模型生成SQL查询\n",
    "    res = llm.invoke(prompt_value).content\n",
    "\n",
    "    # 移除 <SQL> 标签\n",
    "    sql_query = res.replace(\"<SQL>\", \"\").replace(\"</SQL>\", \"\")\n",
    "\n",
    "    return sql_query\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8ab12e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "select red_config from t_scene where scene_id=\"e8f7g6h5-4321-09ab-fedc-ba9876543210\"\n"
     ]
    }
   ],
   "source": [
    "user_prompt = \"查询scene_id='e8f7g6h5-4321-09ab-fedc-ba9876543210'红方配置信息\"\n",
    "generated_sql = gen_system_prompt(user_prompt, model)\n",
    "\n",
    "print(generated_sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d6763e74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'red_config': '{\"units\": [\"early_warning_radar\", \"interceptor\"], \"coverage_area\": \"high_altitude_corridor\"}'}]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 执行查询语句\n",
    "run_sql_query(generated_sql, \"root:123456@localhost:3306/sim_war_game\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e00ef20f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learn_langchain",
   "language": "python",
   "name": "learn_langchain"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
